from diffusers import StableDiffusionXLInpaintPipeline,UNet2DConditionModel
import torch
from PIL import Image,ImageDraw
import numpy as np
import os
scale_index= 1

unet  = UNet2DConditionModel.from_pretrained(
      '/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/sdxl_inpainting_training4/'
      # '/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/sdxl_inpainting_training/checkpoint-14000'
      )
def generate_mask(x1y1x2y2):
      x1 = x1y1x2y2[0]
      y1 = x1y1x2y2[1]
      x2 = x1y1x2y2[2]
      y2 = x1y1x2y2[3]


      mask = np.zeros((360*scale_index,640*scale_index),dtype=np.uint8)
      mask[int(y1):int(y2),int(x1):int(x2)] = 255
      return Image.fromarray(mask).convert('L')

pipeline = StableDiffusionXLInpaintPipeline.from_pretrained(
   '/mnt/afs2d/luotianhang/cache/PretrainedModels/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/462165984030d82259a11f4367a4eed129e94a7b' ,
   variant='fp16'
)

pipeline.unet = unet

pipeline.safety_checker=None
pipeline.to('cuda', dtype=torch.float16)



mem = {}
pad= 100
image = Image.open('/mnt/afs2d01/luotianhang/diffusion_data/process/data_inpainting/inpainting_material/empty_cart_sort/a19/empty_1/frame_15.jpg').convert('RGB')
pad= 100

mem['a furry dog']=[792, 307, 1135, 750]
mem['dog']=[792, 307, 1135, 750]
mem['a black tablet']=[752, 718, 895, 926]
mem['a laptop with an apple on the screen']=[792, 307, 1135, 750]
# mem['computer']=[657, 302, 1080, 754] # 这个词肯能被理解成了集群
mem['cell phone ']=[775, 430, 883, 582]
mem['purple wallet ']=[775, 430, 883, 582]

mem['there is a small dog']=[792, 307, 1175, 750]
mem['there is a small white cat ']=[547, 102, 764, 350]
# mem['a real pet']=[792, 307, 1175, 750]
mem['an aeroplane in the sky']=[547, 102, 764, 350]
mem['a black pig with mouth opened']=[792, 307, 1135, 750]

mem['a black wallet ']=[775, 430, 883, 582]
mem['there is a purse, with money inside ']=[775, 430, 883, 582]
mem['a mobile phone ']=[775, 430, 883, 582]
mem['a brown cat sitting']=[792, 307, 1135, 750]
# mem[' pet ']=[792, 307, 1135, 750]
mem['a yellow dog ']=[792, 307, 1135, 750]
mem['there is a tablet']=[775, 430, 883, 582]

mem['a green bag ']=[547, 102, 764, 350]
mem['there is a safe seat with a pillow on it ']=[757-0, 302-0, 1134+0, 1080+0]
mem['there is a seat ']=[757-0, 302-0, 1134+0, 1080+0]
mem['this is a safeseat ']=[757-0, 302-0, 1134+0, 1080+0]
mem['a safe seat ']=[547-0, 102-0, 764+0, 350+0]
mem['a safeseat ']=[547-0, 102-0, 764+0, 350+0]


mem['big elephant with a long trunk ']=[792, 307, 1135, 750]
mem['a brown bird flying ']=[792, 307, 1135, 750] 
mem['clothes  ']=[792, 307, 1135, 750]
mem['a bread in a sliver plate ']=[792, 307, 1135, 750]
mem[' a birthday cake']=[792, 307, 1135, 750]

width = image.width
height = image.height
count = 0

print('image2image')
for prompt , loc_x1y1x2y2 in mem.items():
   pad = 0
   x1=loc_x1y1x2y2[0]/width*(640*scale_index)
   y1=loc_x1y1x2y2[1]/height*(360*scale_index)
   x2=loc_x1y1x2y2[2]/width*(640*scale_index)
   y2=loc_x1y1x2y2[3]/height*(360*scale_index)

   x1 = max(x1-pad,0)
   y1 = max(y1-pad,0)
   x2 = min(x2+pad,640*scale_index)
   y2 = min(y2+pad,360*scale_index)

   loc_x1y1x2y2=[int(x1),int(y1),int(x2),int(y2)]
   # loc_x1y1x2y2=[0,0,image.width,image.height]
   mask_image = generate_mask(loc_x1y1x2y2)
      
   mask_image.save('./mask_temp.png')
   image = image.resize((640*scale_index,360*scale_index)) 
   print(image.size)
   print(mask_image.size)
   print(loc_x1y1x2y2)
   res_image = pipeline(
         prompt='realistic,best quality,  finely detailed, photorealistic,'+prompt,
         image=image, 
         height=360*scale_index,
         width=640*scale_index,
         mask_image=mask_image,
         num_inference_steps=20,
         negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon",
         ).images[0]
   print(res_image.size)
   res_image = res_image.resize((640*scale_index,360*scale_index)) 
   res_image_draw = ImageDraw.Draw(res_image)
   top_left = (loc_x1y1x2y2[0],loc_x1y1x2y2[1])
   bottom_right = (loc_x1y1x2y2[2],loc_x1y1x2y2[3])
   outline_color = 'green'
   line_width = 2
   res_image_draw.rectangle(
         [top_left,bottom_right],outline=outline_color,width=line_width
   )
   save_dir=os.path.join('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting','demo','sdxl',str(iter),'image2image')
   if not os.path.exists(save_dir):
         os.makedirs(save_dir)

   res_image.save(f"{save_dir}/{str(count)+prompt.replace(' ','_')}.png")
   count+=1

print('text2image')
for prompt , loc_x1y1x2y2 in mem.items():
   pad = 0
   x1=loc_x1y1x2y2[0]/width*(640*scale_index)
   y1=loc_x1y1x2y2[1]/height*(360*scale_index)
   x2=loc_x1y1x2y2[2]/width*(640*scale_index)
   y2=loc_x1y1x2y2[3]/height*(360*scale_index)

   x1 = max(x1-pad,0)
   y1 = max(y1-pad,0)
   x2 = min(x2+pad,640*scale_index)
   y2 = min(y2+pad,360*scale_index)

#    loc_x1y1x2y2=[int(x1),int(y1),int(x2),int(y2)]
   loc_x1y1x2y2=[0,0,640*scale_index,360*scale_index]
   mask_image = generate_mask(loc_x1y1x2y2)
      
   mask_image.save('./mask_temp.png')
   image = image.resize((640*scale_index,360*scale_index))
   res_image = pipeline(
         prompt='realistic,best quality,  finely detailed, photorealistic,'+prompt,
         image=image, 
         mask_image=mask_image,
         num_inference_steps=50,
         negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon",
         ).images[0]
   
   res_image=res_image.resize((640*scale_index,360*scale_index))
   res_image_draw = ImageDraw.Draw(res_image)
   top_left = (loc_x1y1x2y2[0],loc_x1y1x2y2[1])
   bottom_right = (loc_x1y1x2y2[2],loc_x1y1x2y2[3])
   outline_color = 'green'
   line_width = 2
   res_image_draw.rectangle(
         [top_left,bottom_right],outline=outline_color,width=line_width
   )
   save_dir=os.path.join('/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting','demo','sdxl',str(iter),'text2image')
   if not os.path.exists(save_dir):
         os.makedirs(save_dir)

   res_image.save(f"{save_dir}/{str(count)+prompt.replace(' ','_')}.png")
   count+=1